from openai import OpenAI
from langchain_core.messages import HumanMessage,AIMessage

class QWenLLM:
    def __init__(self):
        self.client = OpenAI(
            api_key='sk-965dc39b016c49ecbe29de180f4db2b6',
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        self.system_prompt = '你是一个有用的只能助手'

    def invoke(self,messages:list):
        # 将langchain向小对象转化openai格式
        openai_messages = [
            {"role":"system","content":self.system_prompt}
        ]
        for msg in messages:
            if isinstance(msg,HumanMessage):
                openai_messages.append({"role":"user","content":msg.content})

            if isinstance(msg,AIMessage):
                openai_messages.append({"role":"assistant","content":msg.content})

        completion = self.client.chat.completions.create(
            model='qwen-plus',
            messages=openai_messages
        )
        return AIMessage(content=completion.choices[0].message.content)

if __name__ == '__main__':
    messages = list()
    messages.append(HumanMessage(content="你是谁?"))
    response = QWenLLM().invoke(messages)
    print(response)